import numpy as np
import random
import time
import json
import os
from datetime import datetime
import math

class UAVCommunicationSimulation:
    """无人机通信仿真类，用于测试不同通信协议的性能"""
    
    def __init__(self, num_uavs=20, seed=None, area_size=1000, mobility_model='random_waypoint'):
        """初始化仿真环境
        
        Args:
            num_uavs: 无人机数量
            seed: 随机种子，用于可重复实验
            area_size: 仿真区域大小 (米 x 米)
            mobility_model: 移动模型类型 ('random_waypoint', 'gauss_markov')
        """
        self.num_uavs = num_uavs
        self.area_size = area_size
        self.mobility_model = mobility_model
        
        if seed is not None:
            random.seed(seed)
            np.random.seed(seed)
        
        # 初始化无人机位置、速度和方向
        self.uav_positions = np.random.rand(num_uavs, 3) * np.array([area_size, area_size, 100])  # 3D位置(x,y,高度)
        self.uav_velocities = np.zeros((num_uavs, 3))  # 初始速度为0
        self.uav_max_speeds = np.random.uniform(5, 20, num_uavs)  # 每个无人机的最大速度(5-20 m/s)
        
        # 为随机路点模型创建目标点
        if mobility_model == 'random_waypoint':
            self.waypoints = np.random.rand(num_uavs, 3) * np.array([area_size, area_size, 100])
            self.pause_times = np.zeros(num_uavs)  # 到达目标点后的停留时间
        
        # Gauss-Markov模型参数
        elif mobility_model == 'gauss_markov':
            self.alpha = 0.8  # 时间相关性参数 (0-1)
            self.mean_velocities = np.random.uniform(5, 15, (num_uavs, 3))  # 平均速度
            self.sigma = 5.0  # 高斯噪声标准差
        
        # 协议参数
        self.protocols = {
            'traditional': {
                'range': 300,         # 通信范围 (米)
                'delay_base': 50,     # 基础延迟 (毫秒)
                'delay_var': 20,      # 延迟变化范围 (毫秒)
                'reliability': 0.85,  # 可靠性 (成功率)
                'bandwidth': 250,     # 带宽 (Kbps)
                'emergency_boost': 1.1 # 紧急消息处理提升因子
            },
            'zigbee': {
                'range': 150,
                'delay_base': 40,
                'delay_var': 15,
                'reliability': 0.90,
                'bandwidth': 250,
                'emergency_boost': 1.2
            },
            'wifi_direct': {
                'range': 200,
                'delay_base': 30,
                'delay_var': 10,
                'reliability': 0.92,
                'bandwidth': 1000,
                'emergency_boost': 1.3
            },
            'lora': {
                'range': 800,
                'delay_base': 500,
                'delay_var': 100,
                'reliability': 0.95,
                'bandwidth': 50,
                'emergency_boost': 1.1
            },
            'fanet': {
                'range': 500,
                'delay_base': 60,
                'delay_var': 25,
                'reliability': 0.93,
                'bandwidth': 600,
                'emergency_boost': 1.25
            },
            'dcdp_protocol': {
                'range': 600,         # 通信范围 (米)
                'delay_base': 35,     # 基础延迟 (毫秒)
                'delay_var': 10,      # 延迟变化范围 (毫秒)
                'reliability': 0.97,  # 可靠性 (成功率)
                'bandwidth': 800,     # 带宽 (Kbps)
                'emergency_boost': 1.5, # 紧急消息处理提升因子
                # DCDP特有参数 - 用于三类时延计算
                'processing_delay_factor': 0.05,  # 处理时延因子
                'queueing_delay_min': 5,          # 最小排队时延
                'queueing_delay_max': 25,         # 最大排队时延
                'transmission_delay_factor': 0.8   # 传输时延因子
            }
        }
        
        # 初始化结果统计
        self.reset_statistics()
    
    def reset_statistics(self):
        """重置统计数据"""
        self.statistics = {}
        for protocol in self.protocols:
            self.statistics[protocol] = {
                'total_messages': 0,
                'successful_messages': 0,
                'failed_messages': 0,
                'emergency_messages': 0,
                'emergency_messages_handled': 0,
                'deadline_messages': 0,
                'deadline_met': 0,
                'delays': [],
                'avg_delay': 0
            }
            
            # 为DCDP协议添加三类时延统计
            if protocol == 'dcdp_protocol':
                self.statistics[protocol].update({
                    'propagation_delays': [],   # 信号传播时延
                    'processing_delays': [],    # 数据处理时延
                    'congestion_delays': [],    # 网络拥堵时延
                    'avg_propagation_delay': 0,
                    'avg_processing_delay': 0,
                    'avg_congestion_delay': 0
                })
    
    def update_uav_positions(self, step_time=1.0):
        """更新所有无人机的位置
        
        Args:
            step_time: 时间步长 (秒)
        """
        if self.mobility_model == 'random_waypoint':
            self._update_random_waypoint(step_time)
        elif self.mobility_model == 'gauss_markov':
            self._update_gauss_markov(step_time)
    
    def _update_random_waypoint(self, step_time):
        """随机路点移动模型更新
        
        在这个模型中，无人机会朝着随机目标点移动，到达后会短暂停留，
        然后选择新的目标点继续移动。
        """
        for i in range(self.num_uavs):
            # 如果正在停留状态
            if self.pause_times[i] > 0:
                self.pause_times[i] -= step_time
                continue
            
            # 计算到目标点的距离
            to_waypoint = self.waypoints[i] - self.uav_positions[i]
            distance = np.linalg.norm(to_waypoint)
            
            # 如果到达目标点附近，选择新的目标点并设置停留时间
            if distance < 5.0:  # 5米阈值
                self.waypoints[i] = np.random.rand(3) * np.array([self.area_size, self.area_size, 100])
                self.pause_times[i] = random.uniform(0, 5)  # 0-5秒随机停留
                self.uav_velocities[i] = np.zeros(3)
                continue
            
            # 更新速度方向
            direction = to_waypoint / distance
            speed = min(self.uav_max_speeds[i], distance / step_time)  # 避免超调
            self.uav_velocities[i] = direction * speed
            
            # 更新位置
            self.uav_positions[i] += self.uav_velocities[i] * step_time
            
            # 确保位置在区域范围内
            self.uav_positions[i] = np.clip(
                self.uav_positions[i], 
                [0, 0, 10],  # 最小高度10米
                [self.area_size, self.area_size, 100]  # 最大高度100米
            )
    
    def _update_gauss_markov(self, step_time):
        """Gauss-Markov移动模型更新
        
        在这个模型中，速度具有时间相关性，被建模为一个自回归过程。
        适合模拟平滑的无人机运动。
        """
        for i in range(self.num_uavs):
            # 计算新速度 (Gauss-Markov过程)
            # v(t) = α * v(t-1) + (1-α) * μ + σ * sqrt(1-α^2) * N(0,1)
            rand_normal = np.random.normal(0, 1, 3)
            new_velocity = (self.alpha * self.uav_velocities[i] + 
                          (1 - self.alpha) * self.mean_velocities[i] +
                          self.sigma * np.sqrt(1 - self.alpha**2) * rand_normal)
            
            # 确保速度不超过最大速度
            speed = np.linalg.norm(new_velocity)
            if speed > self.uav_max_speeds[i]:
                new_velocity = new_velocity / speed * self.uav_max_speeds[i]
            
            self.uav_velocities[i] = new_velocity
            
            # 更新位置
            self.uav_positions[i] += self.uav_velocities[i] * step_time
            
            # 边界检查 - 如果到达边界，反射速度
            for j in range(3):
                if self.uav_positions[i, j] < 0:
                    self.uav_positions[i, j] = -self.uav_positions[i, j]
                    self.uav_velocities[i, j] = -self.uav_velocities[i, j]
                elif j < 2 and self.uav_positions[i, j] > self.area_size:  # x, y坐标
                    self.uav_positions[i, j] = 2 * self.area_size - self.uav_positions[i, j]
                    self.uav_velocities[i, j] = -self.uav_velocities[i, j]
                elif j == 2 and self.uav_positions[i, j] > 100:  # z坐标(高度)
                    self.uav_positions[i, j] = 200 - self.uav_positions[i, j]
                    self.uav_velocities[i, j] = -self.uav_velocities[i, j]
                elif j == 2 and self.uav_positions[i, j] < 10:  # 最小高度
                    self.uav_positions[i, j] = 20 - self.uav_positions[i, j]
                    self.uav_velocities[i, j] = -self.uav_velocities[i, j]
    
    def calculate_distance(self, uav1_idx, uav2_idx):
        """计算两架无人机之间的距离
        
        Args:
            uav1_idx: 第一架无人机的索引
            uav2_idx: 第二架无人机的索引
            
        Returns:
            两架无人机之间的欧氏距离
        """
        pos1 = self.uav_positions[uav1_idx]
        pos2 = self.uav_positions[uav2_idx]
        return np.linalg.norm(pos1 - pos2)
    
    def is_in_range(self, uav1_idx, uav2_idx, protocol):
        """判断两架无人机是否在通信范围内
        
        Args:
            uav1_idx: 第一架无人机的索引
            uav2_idx: 第二架无人机的索引
            protocol: 使用的通信协议
            
        Returns:
            如果两架无人机在通信范围内则返回True，否则返回False
        """
        distance = self.calculate_distance(uav1_idx, uav2_idx)
        return distance <= self.protocols[protocol]['range']
    
    def calculate_delay(self, uav1_idx, uav2_idx, protocol, message_size=1, is_emergency=False):
        """计算消息传输延迟
        
        Args:
            uav1_idx: 发送无人机的索引
            uav2_idx: 接收无人机的索引
            protocol: 使用的通信协议
            message_size: 消息大小 (KB)
            is_emergency: 是否为紧急消息
            
        Returns:
            消息传输延迟 (毫秒)
        """
        distance = self.calculate_distance(uav1_idx, uav2_idx)
        protocol_params = self.protocols[protocol]
        
        # 基础延迟 + 距离因子 + 随机变化
        delay = protocol_params['delay_base']
        delay += (distance / protocol_params['range']) * protocol_params['delay_base'] * 0.5
        delay += random.uniform(-protocol_params['delay_var'], protocol_params['delay_var'])
        
        # 考虑消息大小和带宽
        delay += (message_size * 8 / protocol_params['bandwidth']) * 1000
        
        # 如果是紧急消息，则减少延迟
        if is_emergency:
            delay /= protocol_params['emergency_boost']
        
        return max(1, delay)  # 确保延迟至少为1毫秒
    
    def calculate_dcdp_delays(self, uav1_idx, uav2_idx, message_size, is_emergency):
        """计算DCDP协议的三类时延
        
        根据专利说明书，DCDP协议考虑的三类时延为：
        1. 信号传播时延 (Propagation Delay): 与距离直接相关
        2. 数据处理时延 (Processing Delay): 与消息大小和UAV处理能力相关
        3. 网络拥堵时延 (Congestion Delay): 基于网络状态的预测
        
        参数:
            uav1_idx: 发送无人机的索引
            uav2_idx: 接收无人机的索引
            message_size: 消息大小 (KB)
            is_emergency: 是否为紧急消息
            
        返回:
            (传播时延, 处理时延, 拥堵时延, 总时延) 元组，单位：毫秒
        """
        protocol_params = self.protocols['dcdp_protocol']
        distance = self.calculate_distance(uav1_idx, uav2_idx)
        
        # 1. 信号传播时延 - 与距离直接相关
        # 专利公式: T_prop = d/v, 其中d是距离，v是信号传播速度(光速为3*10^8 m/s)
        signal_speed = 3.0e8  # 米/秒
        propagation_delay = (distance / signal_speed) * 1000  # 转换为毫秒
        
        # 2. 数据处理时延 - 与消息大小和处理能力相关
        # 专利公式: T_proc = S / R, 其中S是消息大小，R是处理速率
        processing_rate = 10  # MB/s (假设值)
        processing_delay = (message_size / (processing_rate * 1024)) * 1000  # 转换为毫秒
        
        # 3. 网络拥堵时延 - 基于网络负载和拥塞窗口的预测
        # 专利公式: T_cong = f(load, window), 负载和窗口的函数
        # 简化实现：基于网络负载因子和链路质量
        
        # 随机生成网络负载因子 (0.1-1.0)
        network_load = random.uniform(0.1, 1.0)
        
        # 距离越远，链路质量越差，拥塞几率越高
        link_quality = max(0.1, 1.0 - (distance / protocol_params['range']) * 0.8)
        
        # 拥塞预测基本公式
        # 基础拥塞时延根据链路质量和负载调整
        base_congestion = protocol_params['queueing_delay_max'] * network_load
        congestion_factor = 1.0 / link_quality
        congestion_delay = base_congestion * congestion_factor
        
        # 应用拥塞控制算法 (DCDP专有)
        # 当拥塞高于阈值时，根据专利描述采用自适应控制
        congestion_threshold = protocol_params['queueing_delay_max'] * 0.7
        if congestion_delay > congestion_threshold:
            # 自适应调整，降低拥塞时延，实现拥塞控制
            adaptation_factor = 0.7
            congestion_delay *= adaptation_factor
        
        # 紧急消息优化处理 - 根据专利说明使用优先级调整
        if is_emergency:
            emergency_factor = protocol_params['emergency_boost']
            # 紧急消息的处理时延和拥塞时延会降低
            processing_delay /= emergency_factor
            congestion_delay /= emergency_factor
            # 注意：传播时延不受影响，因为它是物理限制
        
        # 总延迟
        total_delay = propagation_delay + processing_delay + congestion_delay
        
        return propagation_delay, processing_delay, congestion_delay, total_delay
    
    def is_transmission_successful(self, uav1_idx, uav2_idx, protocol, is_emergency=False):
        """判断消息传输是否成功
        
        Args:
            uav1_idx: 发送无人机的索引
            uav2_idx: 接收无人机的索引
            protocol: 使用的通信协议
            is_emergency: 是否为紧急消息
            
        Returns:
            如果消息传输成功则返回True，否则返回False
        """
        # 首先检查是否在范围内
        if not self.is_in_range(uav1_idx, uav2_idx, protocol):
            return False
        
        protocol_params = self.protocols[protocol]
        reliability = protocol_params['reliability']
        
        # 紧急消息有更高的成功率
        if is_emergency:
            reliability = min(0.99, reliability * protocol_params['emergency_boost'])
        
        # 距离越远，成功率越低
        distance = self.calculate_distance(uav1_idx, uav2_idx)
        distance_factor = 1 - (distance / protocol_params['range']) * 0.3
        effective_reliability = reliability * distance_factor
        
        return random.random() < effective_reliability
    
    def simulate_message(self, protocol, is_emergency=False, has_deadline=False, deadline=100):
        """模拟单条消息的传输
        
        Args:
            protocol: 使用的通信协议
            is_emergency: 是否为紧急消息
            has_deadline: 是否有截止时间
            deadline: 截止时间 (毫秒)
            
        Returns:
            (是否成功, 延迟, 是否满足截止时间, 三类时延数据)
        """
        # 随机选择发送和接收无人机
        sender = random.randint(0, self.num_uavs-1)
        receiver = random.randint(0, self.num_uavs-1)
        while receiver == sender:
            receiver = random.randint(0, self.num_uavs-1)
        
        # 生成随机消息大小 (0.5-5KB)
        message_size = random.uniform(0.5, 5)
        
        # 计算传输是否成功
        success = self.is_transmission_successful(sender, receiver, protocol, is_emergency)
        
        # 计算延迟
        if protocol == 'dcdp_protocol':
            # 使用三类时延模型
            prop_delay, proc_delay, cong_delay, delay = self.calculate_dcdp_delays(
                sender, receiver, message_size, is_emergency)
            delay_components = (prop_delay, proc_delay, cong_delay)
        else:
            # 使用基础延迟模型
            delay = self.calculate_delay(sender, receiver, protocol, message_size, is_emergency) if success else float('inf')
            delay_components = None
        
        # 检查是否满足截止时间
        deadline_met = delay <= deadline if has_deadline and success else None
        
        return success, delay, deadline_met, delay_components
    
    def run_simulation(self, num_rounds=1000, emergency_rate=0.1, deadline_rate=0.3, mobility_update_interval=10):
        """运行仿真
        
        Args:
            num_rounds: 仿真轮次
            emergency_rate: 紧急消息的比例
            deadline_rate: 带有截止时间消息的比例
            mobility_update_interval: 多少轮更新一次无人机位置
            
        Returns:
            仿真统计结果
        """
        self.reset_statistics()
        start_time = time.time()
        
        # 仿真中产生的三类时延数据跟踪
        dcdp_delay_components = {
            'propagation': [],
            'processing': [],
            'congestion': []
        }
        
        # 用于跟踪时间序列的时延数据
        delays_over_time = {
            'time_points': [],
            'propagation': [],
            'processing': [],
            'congestion': []
        }
        
        # 按距离分组的时延数据
        distance_groups = ['0-100', '100-200', '200-300', '300-400', '400-500', '>500']
        delays_by_distance = {
            'groups': distance_groups,
            'propagation': [[] for _ in range(len(distance_groups))],
            'processing': [[] for _ in range(len(distance_groups))],
            'congestion': [[] for _ in range(len(distance_groups))]
        }
        
        # 每100轮统计一次时间序列数据
        time_interval = 100
        
        for round_num in range(num_rounds):
            # 定期更新无人机位置 - 模拟移动性
            if round_num % mobility_update_interval == 0:
                self.update_uav_positions(step_time=1.0)
            
            # 收集时间序列数据
            if round_num % time_interval == 0:
                delays_over_time['time_points'].append(round_num)
                # 占位符，稍后填充
                delays_over_time['propagation'].append(0)
                delays_over_time['processing'].append(0)
                delays_over_time['congestion'].append(0)
                # 计数器和累计值
                count = 0
                prop_sum = proc_sum = cong_sum = 0
            
            # 为每个协议模拟一条消息
            for protocol in self.protocols:
                # 确定消息类型
                is_emergency = random.random() < emergency_rate
                has_deadline = random.random() < deadline_rate
                deadline = random.uniform(50, 200) if has_deadline else None
                
                # 模拟消息传输
                success, delay, deadline_met, delay_components = self.simulate_message(
                    protocol, is_emergency, has_deadline, deadline)
                
                # 更新统计信息
                stats = self.statistics[protocol]
                stats['total_messages'] += 1
                
                if success:
                    stats['successful_messages'] += 1
                    stats['delays'].append(delay)
                else:
                    stats['failed_messages'] += 1
                
                if is_emergency:
                    stats['emergency_messages'] += 1
                    if success:
                        stats['emergency_messages_handled'] += 1
                
                if has_deadline:
                    stats['deadline_messages'] += 1
                    if deadline_met:
                        stats['deadline_met'] += 1
                
                # 记录DCDP协议的三类时延
                if protocol == 'dcdp_protocol' and success and delay_components:
                    prop_delay, proc_delay, cong_delay = delay_components
                    stats['propagation_delays'].append(prop_delay)
                    stats['processing_delays'].append(proc_delay)
                    stats['congestion_delays'].append(cong_delay)
                    
                    # 也添加到全局跟踪
                    dcdp_delay_components['propagation'].append(prop_delay)
                    dcdp_delay_components['processing'].append(proc_delay)
                    dcdp_delay_components['congestion'].append(cong_delay)
                    
                    # 为时间序列数据累计
                    if round_num % time_interval == 0:
                        count += 1
                        prop_sum += prop_delay
                        proc_sum += proc_delay
                        cong_sum += cong_delay
                    
                    # 按距离分组记录时延
                    if protocol == 'dcdp_protocol' and success and len(delay_components) == 3:
                        # 随机选择发送和接收无人机（简化模拟）
                        sender = random.randint(0, self.num_uavs-1)
                        receiver = random.randint(0, self.num_uavs-1)
                        while receiver == sender:
                            receiver = random.randint(0, self.num_uavs-1)
                        
                        # 计算距离并确定分组
                        distance = self.calculate_distance(sender, receiver)
                        group_idx = min(int(distance / 100), len(distance_groups) - 1)
                        
                        # 添加到相应分组
                        delays_by_distance['propagation'][group_idx].append(prop_delay)
                        delays_by_distance['processing'][group_idx].append(proc_delay)
                        delays_by_distance['congestion'][group_idx].append(cong_delay)
            
            # 更新时间序列数据
            if round_num % time_interval == 0 and count > 0:
                idx = round_num // time_interval
                delays_over_time['propagation'][idx] = prop_sum / count
                delays_over_time['processing'][idx] = proc_sum / count
                delays_over_time['congestion'][idx] = cong_sum / count
        
        # 计算平均延迟
        for protocol in self.protocols:
            stats = self.statistics[protocol]
            if stats['delays']:
                stats['avg_delay'] = sum(stats['delays']) / len(stats['delays'])
            
            # 计算DCDP的三类时延平均值
            if protocol == 'dcdp_protocol':
                if stats['propagation_delays']:
                    stats['avg_propagation_delay'] = sum(stats['propagation_delays']) / len(stats['propagation_delays'])
                if stats['processing_delays']:
                    stats['avg_processing_delay'] = sum(stats['processing_delays']) / len(stats['processing_delays'])
                if stats['congestion_delays']:
                    stats['avg_congestion_delay'] = sum(stats['congestion_delays']) / len(stats['congestion_delays'])
        
        # 处理按距离分组的时延数据
        avg_delays_by_distance = {
            'distance_groups': distance_groups,
            'propagation_delays': [],
            'processing_delays': [],
            'congestion_delays': []
        }
        
        for i in range(len(distance_groups)):
            # 计算每个距离组的平均时延
            prop_delays = delays_by_distance['propagation'][i]
            proc_delays = delays_by_distance['processing'][i]
            cong_delays = delays_by_distance['congestion'][i]
            
            avg_delays_by_distance['propagation_delays'].append(
                sum(prop_delays) / max(1, len(prop_delays)))
            avg_delays_by_distance['processing_delays'].append(
                sum(proc_delays) / max(1, len(proc_delays)))
            avg_delays_by_distance['congestion_delays'].append(
                sum(cong_delays) / max(1, len(cong_delays)))
        
        simulation_duration = time.time() - start_time
        
        # 简化协议参数，只保留带宽参数用于绘图
        simplified_protocol_params = {}
        for protocol, params in self.protocols.items():
            simplified_protocol_params[protocol] = {
                'bandwidth': params['bandwidth'],
                'range': params['range']
            }
        
        return {
            'statistics': self.statistics,
            'simulation_params': {
                'num_uavs': self.num_uavs,
                'num_rounds': num_rounds,
                'emergency_rate': emergency_rate,
                'deadline_rate': deadline_rate,
                'duration': simulation_duration,
                'protocol_params': simplified_protocol_params,
                'mobility_model': self.mobility_model,
                'area_size': self.area_size
            },
            'dcdp_delay_components': dcdp_delay_components,
            'delays_over_time': delays_over_time,
            'delays_by_distance': avg_delays_by_distance,
            'timestamp': datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        }
    
    def save_results(self, results, filename='simulation_results.json'):
        """保存仿真结果到JSON文件
        
        Args:
            results: 仿真结果
            filename: 输出文件名
        """
        # 拷贝结果，避免修改原始数据
        json_safe_results = self._convert_to_json_serializable(results)
        
        os.makedirs(os.path.dirname(filename), exist_ok=True)
        with open(filename, 'w') as f:
            json.dump(json_safe_results, f, indent=4)
        print(f"仿真结果已保存至 {filename}")
    
    def _convert_to_json_serializable(self, obj):
        """将对象转换为JSON可序列化的格式
        
        Args:
            obj: 需要转换的对象
            
        Returns:
            JSON可序列化的对象
        """
        if isinstance(obj, np.ndarray):
            return obj.tolist()
        elif isinstance(obj, dict):
            return {k: self._convert_to_json_serializable(v) for k, v in obj.items()}
        elif isinstance(obj, list):
            return [self._convert_to_json_serializable(item) for item in obj]
        elif isinstance(obj, tuple):
            return [self._convert_to_json_serializable(item) for item in obj]
        elif isinstance(obj, np.integer):
            return int(obj)
        elif isinstance(obj, np.floating):
            return float(obj)
        elif isinstance(obj, np.bool_):
            return bool(obj)
        else:
            return obj 